import torch.nn as nn
import numpy as np
import torch
import torch.nn.functional as F
from method.base import mlp, DiffAttention
from envs.model.utils import *
from configs.config import BaseEnvConfig
from .helpers import (
    cosine_beta_schedule,
    linear_beta_schedule,
    vp_beta_schedule,
    extract,
    Losses
)

class DiffStatePredictor(nn.Module):
    def __init__(self, config, 
                       graph_model, 
                       device,
                       beta_schedule='linear',
                       n_timesteps=100,
                       loss_type='l1',
                       clip_denoised=False,
                       predict_epsilon=True):
        """
        This function predicts the next state given the current state as input.
        It uses a graph model to encode the state into a latent space and predict each human's next state.
        """
        super().__init__()
        self.trainable = True
        self.graph_model = graph_model
        self.human_motion_predictor = mlp(config.gcn.X_dim, config.model_predictive_rl.motion_predictor_dims)
        self.tmp_config = BaseEnvConfig()
        self.device = device

        self.embedding_dim = config.gcn.X_dim

        self.action_space = None
        self.action_space_dim = 2

        if self.action_space is None:
            self.action_space = build_action_space()

        graph_dims = config.gcn.wr_dims
        self.model = DiffAttention(self.embedding_dim, self.action_space_dim, graph_dims)

        if beta_schedule == 'linear':
            betas = linear_beta_schedule(n_timesteps)
        elif beta_schedule == 'cosine':
            betas = cosine_beta_schedule(n_timesteps)
        elif beta_schedule == 'vp':
            betas = vp_beta_schedule(n_timesteps)

        alphas = 1. - betas
        alphas_cumprod = torch.cumprod(alphas, axis=0)
        alphas_cumprod_prev = torch.cat([torch.ones(1), alphas_cumprod[:-1]])

        self.n_timesteps = int(n_timesteps)
        self.clip_denoised = clip_denoised
        self.predict_epsilon = predict_epsilon

        self.register_buffer('betas', betas)
        self.register_buffer('alphas_cumprod', alphas_cumprod)
        self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev)

        # calculations for diffusion q(x_t | x_{t-1}) and others
        self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod))
        self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - alphas_cumprod))
        self.register_buffer('log_one_minus_alphas_cumprod', torch.log(1. - alphas_cumprod))
        self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod))
        self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alphas_cumprod - 1))

        # calculations for posterior q(x_{t-1} | x_t, x_0)
        posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
        self.register_buffer('posterior_variance', posterior_variance)

        # log calculation clipped because the posterior variance
        # is 0 at the beginning of the diffusion chain
        self.register_buffer('posterior_log_variance_clipped',
                             torch.log(torch.clamp(posterior_variance, min=1e-20)))
        self.register_buffer('posterior_mean_coef1',
                             betas * np.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))
        self.register_buffer('posterior_mean_coef2',
                             (1. - alphas_cumprod_prev) * np.sqrt(alphas) / (1. - alphas_cumprod))

        self.loss_fn = Losses[loss_type]()

    def predict_start_from_noise(self, x_t, t, noise):
        """
            if self.predict_epsilon, model output is (scaled) noise;
            otherwise, model predicts x0 directly
        """
        if self.predict_epsilon:
            return (
                    extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
                    extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * noise
            )
        else:
            return noise
        
    def q_posterior(self, x_start, x_t, t):
        posterior_mean = (
                extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
                extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
        )
        posterior_variance = extract(self.posterior_variance, t, x_t.shape)
        posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
        return posterior_mean, posterior_variance, posterior_log_variance_clipped
    
    @torch.no_grad()
    def p_mean_variance(self, x, t, s):
        x_recon = self.predict_start_from_noise(x, t=t, noise=self.model(x, t, s)) # [batch, drone_num, dim]
        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(
            x_start=x_recon, x_t=x, t=t)
        return model_mean, posterior_variance, posterior_log_variance

    @torch.no_grad()
    def p_sample(self, x, t, s):
        b = x.shape[0]
        model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, s=s)
        noise = torch.rand_like(x)
        nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
        return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise

    @torch.no_grad()
    def p_sample_loop(self, state, shape):
        # device = self.betas.device
        device = self.device
        batch_size = shape[0]
        x = torch.randn(shape, device=device) # noise
        for i in reversed(range(0, self.n_timesteps)):
            timesteps = torch.full((batch_size,), i, device=device, dtype=torch.long)
            x = self.p_sample(x, timesteps, state)  # denoise
        return x

    @torch.no_grad()
    def sample(self, state):
        """
        state: [batch, drone+human, dim]
        """
        batch_size, drone_num = state.shape[0], self.tmp_config.env.robot_num
        ## 潜在向量空间
        shape = (batch_size, drone_num, self.action_space_dim)

        action = self.p_sample_loop(state, shape)
        # action = self.decoder(action_feat)  # [batch, drone_num, 2]
        # return F.softmax(action, dim=-1)
        return action
    
    def q_sample(self, x_start, t, noise=None):
        if noise is None:
            noise = torch.randn_like(x_start)

        sample = (
                extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
                extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
        )
        return sample
    
    def p_losses(self, next_state, state, t, weights=1.0):
        noise = torch.randn_like(next_state)
        x_noisy = self.q_sample(x_start=next_state, t=t, noise=noise)
        x_recon = self.model(x_noisy, t, state)
        assert noise.shape == x_recon.shape
        if self.predict_epsilon:
            return self.loss_fn(x_recon, noise, weights)
        else:
            return self.loss_fn(x_recon, next_state, weights)

    def loss(self, next_drone_states, state, weights=1.0):
        batch_size = state[0].shape[0]
        assert len(state[0].shape) == 3
        assert len(state[1].shape) == 3

        state_embedding = self.graph_model(state).detach()
        next_human_states = self.human_motion_predictor(state_embedding)[:, self.tmp_config.env.robot_num:,:]

        t = torch.randint(0, self.n_timesteps, (batch_size,), device=self.device)
        return self.p_losses(next_drone_states, state_embedding, t, weights), next_human_states

    def forward(self, state, detach=False):
        """ Predict the next state tensor given current state as input.

        :return: tensor of shape (batch_size, # of agents, feature_size)
        """
        assert len(state[0].shape) == 3   # [batch, human_num, 4]
        assert len(state[1].shape) == 3   # [batch, drone_num, 4]

        state_embedding = self.graph_model(state)  # [batch, drone+human, 32]
        if detach:
            state_embedding = state_embedding.detach()
        # diffusion predict(inference)
        next_robot_state = self.sample(state_embedding) # [batch, drone_num, 2]
        # _, ind = torch.max(next_robot_logit, dim=2)
        
        # next_robot_state = self.action_space[ind]
        # predict
        next_human_states = self.human_motion_predictor(state_embedding)[:, self.tmp_config.env.robot_num:,:]

        next_observation = [next_robot_state, next_human_states]
        return next_observation